"""
@Time: 2020/12/10 上午 9:23
@Author: jinzhuan
@File: fn_test.py
@Desc: 
"""
import sys
sys.path.append('/data/zhuoran/code/cognlp')
from cognlp import *
from cognlp.io.loader.fn import FrameNetLoader
from cognlp.io.processor.fn import FrameNetProcessor
from cognlp.utils.util import save_json, load_json
from cognlp.core.dataset import FrameNetDataset
import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.utils.data import RandomSampler
from cognlp.core.metrics import SpanFPreRecMetric
from cognlp.core.trainer import Trainer
from cognlp.core.dataset import NerDataset
from cognlp.models.ner.bert_ner import Bert4Ner
from cognlp.config.ner.conll2003 import args
import logging

device = torch.device('cuda')
torch.cuda.set_device(0)
vocabulary = Vocabulary.load(filepath="../data/fn/framenet/data/vocabulary.txt")
train_table = DataTable.load_table(path='../data/fn/framenet/data/train.json')
dev_table = DataTable.load_table(path='../data/fn/framenet/data/dev.json')
train_data = FrameNetDataset(train_table, device=device)
dev_data = FrameNetDataset(dev_table, device=device)
train_sampler = RandomSampler(train_data)
dev_sampler = RandomSampler(dev_data)
model = Bert4Fn(len(vocabulary))

loss = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.00001)
metric = AccuracyMetric()

trainer = Trainer(model=model, train_data=train_data, dev_data=dev_data, n_epochs=20, batch_size=64,
                 loss=loss, optimizer=optimizer, scheduler=None, metrics=metric,
                 train_sampler=train_sampler, dev_sampler=dev_sampler, drop_last=False,
                 gradient_accumulation_steps=1, num_workers=0, print_every=None,
                 save_path="../data/fn/framenet/model", save_file=None, validate_steps=None, save_steps=None,
                 scheduler_steps=None, grad_norm=None, use_tqdm=True, device=device, device_ids=[0, 3, 4, 5],
                 callbacks=None, check_code_level=logging.INFO, metric_key=None, writer_path="../data/fn/framenet/tensorboard",
                 fp16=False, fp16_opt_level='O1', seed=527, checkpoint_path=None,)
trainer.train()
